import os
from jinja2 import Environment, FileSystemLoader, select_autoescape
from langchain_core.prompts import PromptTemplate

from ..config.annotation import cache
from ..prompts.common import prompt_context
from ..config.fuse import get_langfuse_prompt
from ..utils.str import render

jinja_env = Environment(
    loader=FileSystemLoader(os.path.join(os.path.dirname(__file__), "tpl")),
    autoescape=select_autoescape(),
    trim_blocks=True,
    lstrip_blocks=True,
)
os_env = Environment(
    loader=FileSystemLoader("/"),  # 使用根目录作为基础路径
    autoescape=select_autoescape(),
    trim_blocks=True,
    lstrip_blocks=True,
)


@cache()
def get_prompt_template(prompt_name: str, language: str = None, type: str = None, **kwargs) -> str:
    """
    Load and return a prompt template using Jinja2.

    Args:
        prompt_name: Name of the prompt template file (without .md extension)
        language: Language code (e.g., 'en', 'cn'). If provided, will try to load language-specific template first
        type: Type code. If provided, will try to load type-specific template first

    Returns:
        The template string with proper variable substitution syntax
    """

    def get_template_name_with_type_and_language(base_name: str, type_code: str = None, lang: str = None) -> str:
        """Generate template name with type and language suffix"""
        parts = [base_name]
        if type_code:
            parts.append(type_code)
        if lang:
            parts.append(lang)
        return f"{'_'.join(parts)}.md"

    def template_exists(template_name: str, env: Environment) -> bool:
        """Check if template exists in the given environment"""
        try:
            env.get_template(template_name)
            return True
        except:
            return False

    def try_template_variants(base_name: str, type_code: str = None, lang: str = None):
        """Try different template variants in priority order"""
        variants = []

        # 生成优先级顺序的模板名称
        if type_code and type_code.strip() and lang and lang.strip():
            variants.append((get_template_name_with_type_and_language(base_name, type_code, lang),
                             f"{base_name}_{type_code}_{lang}"))
        if type_code and type_code.strip():
            variants.append(
                (get_template_name_with_type_and_language(base_name, type_code), f"{base_name}_{type_code}"))
        if lang and lang.strip():
            variants.append((get_template_name_with_type_and_language(base_name, None, lang), f"{base_name}_{lang}"))
        variants.append((f"{base_name}.md", base_name))

        return variants

    try:
        # 获取环境变量中的模板路径
        agent_tpl_path = os.getenv('AGENT_TPL_PATH')

        # 优先检查环境变量指定的路径
        if agent_tpl_path:
            for template_name, _ in try_template_variants(prompt_name, type, language):
                custom_template_path = os.path.join(agent_tpl_path, template_name)
                if os.path.exists(custom_template_path):
                    template = os_env.get_template(custom_template_path)
                    return template.render(**kwargs)

        # 检查 langfuse 模板
        for _, langfuse_name in try_template_variants(prompt_name, type, language):
            if langfuse_template := get_langfuse_prompt(langfuse_name):
                return render(langfuse_template, **kwargs)

        # 检查内置模板
        for template_name, _ in try_template_variants(prompt_name, type, language):
            if template_exists(template_name, jinja_env):
                template = jinja_env.get_template(template_name)
                return template.render(**kwargs)

        # 如果所有变体都不存在，使用默认模板
        template = jinja_env.get_template(f"{prompt_name}.md")
        return template.render(**kwargs)
    except Exception as e:
        raise ValueError(f"Error loading template {prompt_name}: {e}")


def get_prompt_context(context: str) -> str:
    return f"""
<context>
{PromptTemplate.from_template(template=prompt_context).format(context=context)}
</context>
"""
